# -*- coding: utf-8 -*-
# =============================================================================
#     FileName:
#         Desc:
#       Author: GGA
#        Email:
#     HomePage:
#      Version: 1.0.1
#   LastChange: 2020-12-20
#      History:
# =============================================================================

import datetime
import logging
import os.path
import traceback

import fire
import redis

from utils.logger_helper import LoggerHelper

app_name = "redis_big_keys"
app_version = "1.0"
LoggerHelper.init_logger(logger_name=app_name, logger_level=logging.INFO)
logger = LoggerHelper.get_logger(logger_name=app_name)


class RedisSample(object):
    def __init__(self, redis_host, redis_port, redis_auth,
                 key_pattern="*", sample_key_count=10000, start_pos=0):
        self.redis_host = redis_host
        self.redis_port = redis_port
        self.redis_auth = redis_auth
        self.key_pattern = key_pattern
        self.sample_key_count = sample_key_count
        self.start_pos = start_pos
        self.batch_pipe_size = 100
        self.batch_scan_size = 100
        self.redis_conn = redis.StrictRedis(host=self.redis_host, port=self.redis_port, password=self.redis_auth, db=0)

    def _get_idle_seconds(self, key_names):
        key_items = dict()
        if len(key_names) == 0:
            return key_items
        redis_pipe = self.redis_conn.pipeline()
        for key_name in key_names:
            redis_pipe.object("IDLETIME", key=key_name)
        pipe_result = redis_pipe.execute()
        for key_index, key_name in enumerate(key_names):
            key_items[key_name] = pipe_result[key_index]
        return key_items

    def _get_ttl_seconds(self, key_names):
        key_items = dict()
        if len(key_names) == 0:
            return key_items
        redis_pipe = self.redis_conn.pipeline()
        for key_name in key_names:
            redis_pipe.ttl(key_name)
        pipe_result = redis_pipe.execute()
        for key_index, key_name in enumerate(key_names):
            key_items[key_name] = pipe_result[key_index]
        return key_items

    def _get_memory_usage(self, key_names):
        key_items = dict()
        if len(key_names) == 0:
            return key_items
        redis_pipe = self.redis_conn.pipeline()
        for key_name in key_names:
            redis_pipe.memory_usage(key_name)
        pipe_result = redis_pipe.execute()
        for key_index, key_name in enumerate(key_names):
            key_items[key_name] = pipe_result[key_index]
        return key_items

    def _scan_keys(self):
        """
        范围扫描匹配的键值, 并返回匹配的KEYS
        :return: 满足条件的KEYS
        """
        logger.info("开始扫描.")
        start_pos = self.start_pos
        loop_index = 0
        match_keys = []
        while True:
            loop_index += 1
            stop_pos, tmp_match_keys = self.redis_conn.scan(start_pos, self.key_pattern, self.batch_scan_size)
            logger.info("扫描起始点:{},扫描结束点:{},扫描数量:{},匹配Key数量:{}".format(
                start_pos, stop_pos, self.batch_scan_size, len(tmp_match_keys)))
            start_pos = stop_pos
            match_keys += tmp_match_keys
            if len(match_keys) > self.sample_key_count:
                break
        scan_keys = loop_index * self.batch_scan_size
        logger.info("完成扫描,共扫描{}个KEY,获取到{}个KEY".format(
            scan_keys,
            len(match_keys))
        )
        return match_keys, scan_keys

    def _spilt_keys_to_pipe_groups(self, key_names):
        """
        将KEY按照pipe操作批次拆分多个pipe组
        :param key_names:
        :return:
        """
        pipe_groups = []
        group_keys = []
        for key_name in key_names:
            group_keys.append(key_name)
            if len(group_keys) == self.batch_pipe_size:
                pipe_groups.append(group_keys)
                group_keys = []
            pass
        if len(group_keys) > 0:
            pipe_groups.append(group_keys)
        return pipe_groups

    def _get_key_types(self, key_names):
        key_items = dict()
        if len(key_names) == 0:
            return key_items
        redis_pipe = self.redis_conn.pipeline()
        for key_name in key_names:
            redis_pipe.type(key_name)
        pipe_result = redis_pipe.execute()
        for key_index, key_name in enumerate(key_names):
            key_items[key_name] = pipe_result[key_index]
        return key_items

    @classmethod
    def get_expire_seconds(cls, ttl_seconds, idle_seconds):
        if int(ttl_seconds) == -1:
            return 9999999
        return int(ttl_seconds) + int(idle_seconds)

    def _get_sample_keys(self, key_names):
        logger.info("开始获取KEY的基础信息")
        key_info_list = list()
        pip_groups = self._spilt_keys_to_pipe_groups(key_names=key_names)
        group_count = len(pip_groups)
        logger.info("将采样KEY拆分为{}组".format(group_count))
        group_index = 0
        for key_items in pip_groups:
            group_index += 1
            logger.info("开始获取第{}/{}组的KEY信息".format(group_index, group_count))
            memory_items = self._get_memory_usage(key_names=key_items)
            ttl_items = self._get_ttl_seconds(key_names=key_items)
            idle_items = self._get_idle_seconds(key_names=key_items)
            type_items = self._get_key_types(key_names=key_items)
            for key_name in key_items:
                memory_usage = memory_items.get(key_name, 0)
                if memory_usage is None:
                    memory_usage = -1
                ttl_seconds = ttl_items.get(key_name, 0)
                if ttl_seconds is None:
                    ttl_seconds = -1
                idle_seconds = idle_items.get(key_name, 0)
                if idle_seconds is None:
                    idle_seconds = 0
                key_type = type_items.get(key_name, 'NULL')
                if key_type is None:
                    key_type = "NULL"
                expire_seconds = self.get_expire_seconds(ttl_seconds=ttl_seconds, idle_seconds=idle_seconds)
                key_info = {
                    "key_name": key_name,
                    "memory_usage": memory_usage,
                    "ttl_seconds": ttl_seconds,
                    "idle_seconds": idle_seconds,
                    "expire_seconds": expire_seconds,
                    "key_type": key_type
                }
                key_info_list.append(key_info)
        logger.info("完成获取KEY的基础信息")
        return key_info_list

    @classmethod
    def _save_to_excel(cls, key_info_list, sort_type, file_name):
        base_dir = os.path.dirname(os.path.abspath(__file__))
        data_dir = os.path.join(base_dir, "data")
        if not os.path.exists(data_dir):
            os.mkdir(data_dir)
        file_path = os.path.join(data_dir, file_name)
        logger.info("准备将数据导入到文件{}.".format(file_name))
        key_info_list = sorted(key_info_list, key=lambda item: item[sort_type], reverse=True)
        key_lines = list()
        key_line = "|| {:>15} || {:>15} || {:>15} || {:>15} || {:>15} || {} \n".format(
            "key_type",
            "memory_usage",
            "ttl_seconds",
            "idle_seconds",
            "expire_seconds",
            "key_name"
        )
        key_lines.append(key_line)
        for key_info in key_info_list:
            try:
                key_line = "|| {:>15} || {:>15} || {:>15} || {:>15} || {:>15} || {} \n".format(
                    key_info["key_type"].decode("utf-8", errors="ignore"),
                    key_info["memory_usage"],
                    key_info["ttl_seconds"],
                    key_info["idle_seconds"],
                    key_info["expire_seconds"],
                    key_info["key_name"].decode("utf-8", errors="ignore")
                )
                key_lines.append(key_line)
            except Exception as ex:
                logger.warning("报错记录出现异常,异常为:{},堆栈为:{}".format(
                    str(ex), traceback.format_exc()
                ))
        with open(file=file_path, mode="w+") as fh:
            fh.writelines(key_lines)
            fh.flush()
        logger.info("数据导入完成,文件名为: {}".format(file_path))

    def sample_keys(self):
        key_names, scan_keys = self._scan_keys()
        key_info_list = self._get_sample_keys(key_names=key_names)
        # 排序方式,支持："key_type","memory_usage","ttl_seconds","idle_seconds","key_name"
        prefix_file_name = "simple_keys_{}".format(datetime.datetime.now().strftime("%Y%m%d%H%M%S"))
        for sort_type in ["key_type", "memory_usage", "ttl_seconds", "idle_seconds", "key_name", "expire_seconds"]:
            file_name = prefix_file_name + "_by_{}.txt".format(sort_type)
            self._save_to_excel(key_info_list=key_info_list, sort_type=sort_type, file_name=file_name)
        logger.info("扫描完成,共扫描{}个KEY,找到匹配的{}个KEY".format(scan_keys, len(key_names)))


def main(redis_host="127.0.0.1", redis_port=6379, redis_auth="",
         key_pattern="*", sample_key_count=10000, start_position=0):
    """
    :param redis_host: Redis实例地址
    :param redis_port: Redis实例端口
    :param redis_auth: Redis账号秘钥
    :param key_pattern: 需要匹配的KEY正则表达式
    :param sample_key_count: 需要采样的KEY数据
    :param start_position: 采样的开始位点
    :return:
    """
    monitor = RedisSample(redis_host, redis_port, redis_auth,
                          key_pattern, sample_key_count, start_position)
    monitor.sample_keys()


if __name__ == '__main__':
    fire.Fire(main)
